import asyncio
import torch
from datasets import Dataset
from langchain_huggingface import HuggingFaceEmbeddings
from ragas import  evaluate
from ragas.metrics import ContextRelevance, answer_relevancy, faithfulness, ResponseGroundedness

#from retrieve import rephrase_retrieve, get_rag_chain, get_llm, get_retriever
#from retrieve_HyDE import rephrase_retrieve, get_rag_chain, get_llm, get_retriever,llm_HyDE
from retrieve_rrf import rephrase_retrieve, get_rag_chain, get_llm, get_retriever


#这个文件夹，将运用retrieve.py中的方法,调用链条进行检索，获取查询结果,添加ragas评估
chat_history = []



#1、获取向量模型，方法中需要向量模型作为参数传递
embedding_model = HuggingFaceEmbeddings(
    model_name="./bge-base-zh-v1.5",
    model_kwargs={"device": "cuda" if torch.cuda.is_available() else "cpu"},
    encode_kwargs={
        "normalize_embeddings": True
    },  # 输出归一化向量，更适合余弦相似度计算
)


#2、获取大语言模型作为参数传递
llm=get_llm()


#3、创建调用retriever的函数
async def invoke_rag(query,conversation_id,chat_history):
    answer=""
    #3、1创建检索器
    retrieve=get_retriever(k=3,embedding_model=embedding_model)
    #3、2将问题传入进行重塑,将重述后的问题进行rag检索
    retrieve_result=rephrase_retrieve({"query":query,"history":chat_history},llm,retrieve,multi_query_num=5)

    #3、2将rag检索后的结果传入rag链条进行调用，用于大语言模型生成回复
    async for chunk in  get_rag_chain(retrieve_result,llm).astream({"query":query,"history":chat_history}):
       answer +=chunk
       yield chunk
    #3、4将得到的回复结果存进历史对话中
    chat_history.extend(
        [
            {"role":"用户","content":query},
            {"role":"助手","content":answer}
        ]
    )



#4、main函数传输问题，并调用3的函数
if __name__ == '__main__':
    async def main():
        query = ["不动产被侵占了怎么办"]
        for i in query:
            print(f"==={i}===")
            async for chunk in invoke_rag(i,1,chat_history):
                print(chunk,end="",flush=True)
            print('\n')
    asyncio.run(main())

